import sys
import threading
import queue
import pyaudio
import numpy as np
import torch
import time
import re
import os
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess


class ContextManager:
    """上下文管理器 - 用于维护对话历史，提高识别准确率"""

    def __init__(self, max_context_length=5):
        """
        初始化上下文管理器
        :param max_context_length: 最大上下文长度，超过则移除最旧的内容
        """
        self.context_buffer = []  # 存储上下文的缓冲区
        self.max_context_length = max_context_length  # 最大上下文长度

    def update_context(self, text):
        """
        更新上下文缓冲区
        :param text: 新的文本内容
        """
        if text.strip():  # 只有非空文本才添加到上下文
            self.context_buffer.append(text)
            if len(self.context_buffer) > self.max_context_length:
                self.context_buffer.pop(0)  # 移除最旧的内容

    def get_context_prompt(self):
        """
        获取上下文提示 - 用于提供给语音识别模型作为参考
        :return: 最近3句话组成的上下文字符串
        """
        return " ".join(self.context_buffer[-3:])  # 使用最近3句作为上下文

    def clear_context(self):
        """清空上下文"""
        self.context_buffer.clear()


class SemanticSegmenter:
    """语义分割器 - 判断识别结果是否为完整的句子"""

    def __init__(self):
        """
        初始化语义分割器
        定义句子结束符号，包括中英文标点
        """
        self.sentence_endings = ['.', '!', '?', '。', '！', '？', '；', ';']

    def is_complete_sentence(self, text):
        """
        判断是否为完整句子
        :param text: 待判断的文本
        :return: 是否为完整句子
        """
        text = text.strip()
        if not text:
            return False
        # 检查是否以句子结束符结尾
        has_ending = any(text.endswith(ending) for ending in self.sentence_endings)
        # 同时检查长度，如果文本足够长，即使没有结束符也认为是完整句子
        is_long_enough = len(text) > 20
        return has_ending or is_long_enough


class ASRProcessor:
    """
    语音识别处理器 - 负责核心的语音识别功能
    """
    
    def __init__(self, model, context_manager, semantic_segmenter):
        """
        初始化语音识别处理器
        :param model: ASR模型实例
        :param context_manager: 上下文管理器
        :param semantic_segmenter: 语义分割器
        """
        self.model = model
        self.context_manager = context_manager
        self.semantic_segmenter = semantic_segmenter
        self.last_printed_text = ""
        self.waiting_for_complete_sentence = False

    def recognize_speech(self, audio_data):
        """
        识别语音数据
        :param audio_data: 音频数据列表
        :return: 识别结果和状态信息
        """
        # 合并音频数据
        audio_bytes = b''.join(audio_data)
        audio_array = np.frombuffer(audio_bytes, dtype=np.int16)
        audio_float = audio_array.astype(np.float32) / 32768.0

        # 当音频数据足够长时才进行识别
        if len(audio_array) <= 8000:  # 降低最小音频长度要求到0.5秒
            return None, False

        try:
            # 获取上下文提示
            context_prompt = self.context_manager.get_context_prompt()

            res = self.model.generate(
                input=audio_float,
                cache={},  # 缓存参数
                language="auto",  # 自动检测语言
                use_itn=True,  # 使用逆文本标准化
                batch_size_s=60,  # 批处理大小
                merge_vad=True,  # 使用模型内置的VAD功能
                hotword=context_prompt  # 使用上下文作为热词提示
            )

            if res and len(res) > 0:
                # 使用模型原始输出，不进行额外的文本清理
                text = rich_transcription_postprocess(res[0]["text"])
                # 只有在文本非空且与上次输出不同时才返回
                if text.strip() and text != self.last_printed_text:
                    self.last_printed_text = text
                    # 更新上下文
                    self.context_manager.update_context(text)

                    # 检查语义完整性
                    is_complete = self.semantic_segmenter.is_complete_sentence(text)
                    self.waiting_for_complete_sentence = not is_complete
                    
                    return text, is_complete
                else:
                    # 没有识别到有效文本
                    self.waiting_for_complete_sentence = False
                    return None, False
            else:
                # 没有识别结果
                self.waiting_for_complete_sentence = False
                return None, False

        except Exception as e:
            # 出错时重置状态
            self.waiting_for_complete_sentence = False
            raise e

    def should_wait_for_complete_sentence(self):
        """
        是否在等待完整句子
        :return: 布尔值，表示是否在等待完整句子
        """
        return self.waiting_for_complete_sentence

    def reset_state(self):
        """
        重置识别状态
        """
        self.waiting_for_complete_sentence = False


class OptimizedRealTimeASR:
    """
    优化版实时语音识别系统
    主要优化点：
    1. 使用上下文管理提升识别准确率
    2. 实现语义级句子分割，而非简单按时间分割
    3. 增加缓冲区大小，支持更长语音输入
    4. 引入超时机制，确保及时输出结果
    """

    def __init__(self, model_dir="iic/SenseVoiceSmall", model_cache_dir=None):
        """
        初始化实时语音识别系统
        :param model_dir: ASR模型路径
        :param model_cache_dir: 模型缓存目录路径
        """
        # 设置模型缓存目录
        if model_cache_dir:
            # 注意：ModelScope会在指定目录下创建models子目录，所以我们直接使用传入的目录
            os.environ['MODELSCOPE_CACHE'] = model_cache_dir
        
        # 初始化模型，禁用更新检查
        self.model = AutoModel(
            model=model_dir,
            trust_remote_code=True,
            vad_model="iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",  # 使用ModelScope的VAD模型
            vad_kwargs={"max_single_segment_time": 60000},  # 增加单段最大时间到60秒
            device="cuda:0" if torch.cuda.is_available() else "cpu",  # 优先使用GPU
            disable_pbar=True,  # 禁用进度条
            disable_update=True  # 禁用模型更新检查
        )

        # 音频参数配置
        self.chunk_size = 1024  # 音频块大小
        self.format = pyaudio.paInt16  # 音频格式
        self.channels = 1  # 单声道
        self.rate = 16000  # 采样率 16kHz
        self.frame_duration_ms = 30  # 每帧时长（毫秒）
        self.frame_size = int(self.rate * self.frame_duration_ms / 1000)  # 每帧样本数

        # 上下文管理
        self.context_manager = ContextManager()

        # 语义分割
        self.semantic_segmenter = SemanticSegmenter()

        # 语音识别处理器
        self.asr_processor = ASRProcessor(self.model, self.context_manager, self.semantic_segmenter)

        # 音频缓冲区
        self.audio_queue = queue.Queue()  # 音频数据队列
        self.speech_buffer = []  # 语音缓冲区
        self.max_speech_buffer_duration = 30  # 增加最大语音缓冲区时长到30秒
        self.max_speech_buffer_frames = int(self.max_speech_buffer_duration * self.rate / self.frame_size)  # 最大帧数

        # 控制标志
        self.recording = False  # 录音状态标志
        self.processing = False  # 处理状态标志
        self.last_speech_time = time.time()  # 最后一次检测到语音的时间（用于超时机制）

    def audio_callback(self, in_data, frame_count, time_info, status):
        """
        音频数据回调函数 - PyAudio流式录音的回调
        :param in_data: 输入的音频数据
        :param frame_count: 帧数量
        :param time_info: 时间信息
        :param status: 状态标志
        :return: 返回音频数据和继续标志
        """
        self.audio_queue.put(in_data)  # 将音频数据放入队列
        return (in_data, pyaudio.paContinue)

    def start_recording(self):
        """开始录音"""
        if self.recording:
            return

        self.recording = True
        self.audio = pyaudio.PyAudio()  # 创建PyAudio对象

        # 打开音频流
        self.stream = self.audio.open(
            format=self.format,
            channels=self.channels,
            rate=self.rate,
            input=True,
            frames_per_buffer=self.chunk_size,
            stream_callback=self.audio_callback  # 设置回调函数
        )

    def stop_recording(self):
        """停止录音"""
        self.recording = False
        if hasattr(self, 'stream'):
            self.stream.stop_stream()
            self.stream.close()
        if hasattr(self, 'audio'):
            self.audio.terminate()

    def process_audio(self):
        """处理音频数据 - 核心处理逻辑"""
        self.processing = True
        speech_frames = []  # 存储当前收集的语音帧
        # 设置超时时间（秒）
        timeout_duration = 10.0  # 增加超时时间到10秒
        last_process_time = time.time()  # 上次处理时间
        processing_delay = 4.0  # 处理间隔增加到4秒

        while self.processing:
            try:
                # 获取音频数据
                if not self.audio_queue.empty():
                    data = self.audio_queue.get()

                    # 收集音频数据
                    speech_frames.append(data)
                    self.last_speech_time = time.time()  # 更新最后语音时间

                    # 防止缓冲区过大
                    if len(speech_frames) > self.max_speech_buffer_frames:
                        speech_frames.pop(0)

                # 定期处理或超时处理
                current_time = time.time()
                # 检查是否需要处理：有语音数据且（到了处理间隔时间或超时了或缓冲区足够大）
                should_process = (
                        len(speech_frames) > 0 and
                        (current_time - last_process_time > processing_delay or  # 定期处理
                         (time.time() - self.last_speech_time) > timeout_duration or  # 超时处理
                         len(speech_frames) > self.max_speech_buffer_frames * 0.8)  # 缓冲区接近满时处理
                )

                if should_process:
                    try:
                        # 进行语音识别
                        result, is_complete = self.asr_processor.recognize_speech(speech_frames)
                        
                        # 如果有识别结果，则打印
                        if result is not None:
                            print(result)
                            
                    except Exception as e:
                        # 出错时也清空缓冲区
                        speech_frames = []
                        self.asr_processor.reset_state()
                        last_process_time = current_time
                        continue
                    finally:
                        # 根据是否在等待完整句子决定是否清空缓冲区
                        if not self.asr_processor.should_wait_for_complete_sentence():
                            # 不需要等待完整句子，清空缓冲区
                            speech_frames = []
                            last_process_time = current_time
                        else:
                            # 在等待完整句子，保留缓冲区数据，只更新处理时间
                            last_process_time = current_time

                    # 短暂休眠避免占用过多CPU
                    time.sleep(0.01)
                else:
                    # 短暂休眠避免占用过多CPU
                    time.sleep(0.01)

            except Exception as e:
                # 出错时也清空缓冲区
                speech_frames = []
                self.asr_processor.reset_state()
                pass

    def run(self):
        """运行实时语音识别"""
        try:
            # 开始录音
            self.start_recording()

            # 在单独的线程中处理音频
            process_thread = threading.Thread(target=self.process_audio)
            process_thread.daemon = True  # 设置为守护线程
            process_thread.start()

            # 等待用户中断
            print("开始实时语音识别，按回车键停止...")
            input()

        except KeyboardInterrupt:
            pass
        finally:
            self.processing = False
            self.stop_recording()


if __name__ == "__main__":
    asr = OptimizedRealTimeASR()
    asr.run()